import time

import yaml


def update_kubelet(ssh_client, cpu_manager_policy, topolofy_manager_policy):
    """使用SSH远程连接更新kubelet配置

    Args:
        ssh_client: paramiko SSHClient实例
        cpu_manager_policy: CPU管理策略
        topolofy_manager_policy: 拓扑管理策略
    """
    allowed_cpu_manager_policies = ["none", "static"]
    if cpu_manager_policy not in allowed_cpu_manager_policies:
        raise ValueError(f"cpuManagerPolicy is not valid: {cpu_manager_policy}, \
                         should be among {allowed_cpu_manager_policies}.")
    allowed_topology_manager_policies = ["none", "best-effort", "restricted", "single-numa-node"]
    if topolofy_manager_policy not in allowed_topology_manager_policies:
        raise ValueError(f"topologyManagerPolicy is not valid: {topolofy_manager_policy}, \
                         should be among {allowed_topology_manager_policies}.")

    config_path = "/var/lib/kubelet/config.yaml"
    cpu_manager_state_path = "/var/lib/kubelet/cpu_manager_state"

    try:
        # 读取远程kubelet config
        stdin, stdout, stderr = ssh_client.exec_command(f"cat {config_path}")
        config_content = stdout.read().decode('utf-8')
        error = stderr.read().decode('utf-8')

        if error:
            raise Exception(f"读取配置文件 {config_path} 失败: {error}")

        # 更新配置
        kubelet_config = yaml.safe_load(config_content)
        if "systemReserved" not in kubelet_config:
            kubelet_config["systemReserved"] = {"cpu": "0.5"}
        if "kubeReserved" not in kubelet_config:
            kubelet_config["kubeReserved"] = {"cpu": "0.5"}
        kubelet_config["cpuManagerPolicy"] = cpu_manager_policy
        kubelet_config["topologyManagerPolicy"] = topolofy_manager_policy
        updated_config = yaml.safe_dump(kubelet_config, default_flow_style=False)

        # 上传新 kubelet config
        sftp = ssh_client.open_sftp()
        try:
            temp_config_path = f"{config_path}.tmp"
            with sftp.file(temp_config_path, 'w') as f:
                f.write(updated_config)
            ssh_client.exec_command(f"mv {temp_config_path} {config_path}")
            ssh_client.exec_command(f"rm -f {cpu_manager_state_path}")

            stdin, stdout, stderr = ssh_client.exec_command("systemctl restart kubelet")
            exit_status = stdout.channel.recv_exit_status()
            
            if exit_status != 0:
                error_msg = stderr.read().decode('utf-8')
                raise Exception(f"重启kubelet服务失败: {error_msg}")

            time.sleep(1)
        finally:
            sftp.close()

    except Exception as e:
        raise Exception(f"更新kubelet配置失败: {str(e)}") from e
